function [p_s, PM_s, pM_ig, RevM_ig, pX_ig, VM_ig, F_r, RW_r] ...   
         = compute_equilibrium_nonlinear( delta_ig, a_star_ig, z_star_ig, a_ig, aM_g, AM_s, AD_s, betaT, beta_s, alphaL_s, alphaI_s, alpha_ks, barL_rs, Z_rs, D, br, ... 
         tau_ig, tau_star_ig, m_ig_0, E_s_0, p_s_0, Dsg, DX, DM, omega_star, sigma, eta, kappa, sigma_star, nu, epsilon, gammaQ, gammaX, gammaM, rhoX, rhoM, tol, adj_inner, adj_outter, ...
         ind_Mshifts_vec, ind_Xshifts_vec, ind_Msample_ig_vec, ind_Xsample_ig_vec)

%This script computes the non-linear predictions of the model for a given set of tariffs
%Corresponding author: Rodrigo Adao
%Date: 09/11/2024

%% Preliminaries

D_s_last = E_s_0;
[R, S] = size(barL_rs);
[N, G] = size(a_ig);

%Inital product imports for correcting parameters
m_g0 = sum( m_ig_0, 1)'; 
%m_s0 = Dsg*m_g0;

%Compute parameters that do not change across iterations

    %terms for variety-level imports
    varphi_ig = ( a_ig.^(1+gammaM*omega_star) ).*( z_star_ig.^( gammaM*(1-sigma) ) ).*( (1 + tau_ig).^( gammaM*(1-sigma) ) );
    varphi_ig = varphi_ig.^( 1/(1+gammaM*omega_star*sigma) ) ;

    tilde_varphiR_ig = ( a_ig.^(1+omega_star) ).*( z_star_ig.^( 1 - gammaM*sigma ) ).*( (1 + tau_ig).^( - gammaM*sigma*(1+omega_star) ) );
    tilde_varphiR_ig = tilde_varphiR_ig.^( 1/(1+gammaM*omega_star*sigma) ) ;
    
    varphiR_ig = tau_ig.*tilde_varphiR_ig;

    %terms for product-level imports
    varphi_g = sum( varphi_ig , 1)';
    varphiR_g = sum( varphiR_ig , 1)';

    mu_g = ( aM_g.^( (1+gammaM*omega_star)/(1+gammaM*omega_star*eta) ) ).*( varphi_g.^( ( (1-eta)*(1+gammaM*omega_star*sigma) )/( (1-sigma)*(1+gammaM*omega_star*eta) ) ) ); 
    mu_gs_aux = ( ( Dsg*m_g0 )'*Dsg )'; %Wlog assign cost of one to products in sectors without impots 
    mu_g(mu_gs_aux<tol) = 1;
    mu_s = Dsg*mu_g;
    
    %terms for sector-level imports
    zeta_s = ( ( AM_s.^(gammaM*omega_star) ).*( mu_s.^( ( 1+gammaM*omega_star*eta )/(1-eta) ) ) ).^(1/(1+gammaM*omega_star*kappa)) ;
    
    %terms for production
    alphaK_s = 1 - alphaL_s - alphaI_s;
    alphaK_rs = repmat(alphaK_s',R,1);
    alphaI_rs = repmat(alphaI_s',R,1);
    alphaL_rs = repmat(alphaL_s',R,1);

    %terms for exports
    deltaX_ig_term = a_star_ig.*( (1 + tau_star_ig).^(-sigma_star) ).*( delta_ig.^(1-sigma_star) ) ;
    deltaX_g_term = sum(deltaX_ig_term,1)';
    deltaX_s = Dsg*deltaX_g_term;
    
    %terms for import tariff revenue
    R_g1 = ( aM_g.^( (1+omega_star)/(1+gammaM*omega_star*eta) ) ).*( varphi_g.^( ( (sigma-eta)*(1+omega_star)  )/( (1-sigma)*(1+gammaM*omega_star*eta) ) ) );
    R_g = R_g1.*varphiR_g;
    R_g(aM_g==0) = 0; %wlog assign zero to products without imports
    R_g(mu_gs_aux<tol) = 0; %wlog assign equal shifters to all products if there are no imports in a sector

    R_s = ( mu_s.^( ( (eta-kappa)*(1+omega_star) )/( (1-eta)*(1+gammaM*omega_star*kappa) ) ) ).*( Dsg*R_g ) ;  
    R_s = DM*R_s;   

    %terms for export tax revenue
    RX_ig_term = a_star_ig.*( (1 + tau_star_ig).^(-sigma_star) ).*( delta_ig.^(1-sigma_star) ).*tau_star_ig ;
    RX_g_term = sum(RX_ig_term,1)';
    RX_s = Dsg*RX_g_term;    
    RX_s = DX*RX_s;

%Compute initial guess of import price index

    PM_s = zeta_s.*( D_s_last.^( gammaM*omega_star/(1+gammaM*omega_star*kappa) ) );
    PM_s(zeta_s==0)=1; %wlog assign price of one to sectors without imports
    %PM_s(m_s0==0)=1; %wlog assign price of one to sectors without imports

%% Solve for prices

% Outter loop: Iterate to solve for sector-level import price
p_s_last = p_s_0;
for j=1:10000
    
    % Inner loop: Iterate to solve for domestic sector-level price (given PM_s)
    p_s=p_s_last;
    for i=1:10000
        %sector-level price index
        P_s = ( AD_s.*( p_s.^(1-kappa) ) + AM_s.*( PM_s.^(1-kappa) ) ).^(1/(1-kappa)) ;
        
        %sector-level input cost
        if nu == 1
            phi_s = exp( alpha_ks'*log( P_s ) );
        else
            phi_s = ( alpha_ks'*( P_s.^(1-nu) ) ).^( 1/(1-nu) );
        end
        
        %sector-region wage   
        w_rs_sectorterm = ( ( alphaL_s.^alphaK_s ).*( p_s ).*( phi_s.^(-alphaI_s) ) ).^(1./(1-alphaI_s)) ;
        w_rs_regionterm = ( ( barL_rs.^(-alphaK_rs) ).*( Z_rs ) ).^(1./(1-alphaI_rs)) ;
        w_rs_regionterm( isnan(w_rs_regionterm) ) = 1; %set wage = 0 for every region-sector without employment
        
        w_rs = w_rs_regionterm.*repmat(w_rs_sectorterm',R,1) ;
        
        %sector-level output value
        Y_s_regionterm = sum( ( Z_rs.*( w_rs.^(-alphaL_rs) ) ).^( 1./alphaK_rs ) ,1)';
        Y_s_sectorterm = ( p_s.*( phi_s.^(-alphaI_s) ) ).^( 1./alphaK_s );
        Y_s = Y_s_sectorterm.*Y_s_regionterm;
        %max(abs( Y_s./sales_s - 1))
        
        %sector-level export value
        X_s = ( deltaX_s ).*( p_s.^(1-sigma_star) );
        %max( abs(X_s./(Dsg*sum(x_ig,1)') - 1))
        
        %Export tax revenue
        RevX = ( RX_s' )*( p_s.^(1-sigma_star) );

        %sector-level final spending shares
        eF_s = ( beta_s.*( P_s.^(1-epsilon) ) )./( beta_s'*( P_s.^(1-epsilon) ) ) ;
        
        %sector-to-sector intermediate spending shares
        eI_ks_sterm = alpha_ks'*( P_s.^(1-nu) );
        eI_ks = diag( P_s.^(1-nu) )*alpha_ks*diag(1./eI_ks_sterm);
        

        %sector-level spending value (final + intermediate)
        tildeE_s = ( eF_s )*(D + RevX + (1-alphaI_s)'*Y_s ) + ( eI_ks*diag(alphaI_s) )*Y_s ;       

        parmR_s = (1 + omega_star)/(1+gammaM*kappa*omega_star);
        R_s_aux = ( ( ( P_s.^(kappa-1) ).*AM_s ).^parmR_s ).*R_s ;
        r_ks = ( repmat( eF_s , 1, S) ).*( repmat( R_s_aux'  , S, 1) );

        E_s_last = ( eye(S)/(eye(S) - r_ks) )*tildeE_s;
        for lcount = 1:1000
            E_s = r_ks*( E_s_last.^parmR_s ) + tildeE_s;

            Edif_s = E_s./E_s_last - 1 ;
            err_Edif_step = max(abs(Edif_s));
            
            if err_Edif_step > tol/100
                E_s_next = E_s_last - adj_outter*(E_s_last - E_s);
                E_s_last = E_s_next;
            else
               break 
            end            
        end    
        
        %sector-level domestic spending value
        ED_s = AD_s.*( (p_s./P_s).^(1-kappa) ).*E_s;
        D_s = E_s.*( P_s.^(kappa-1) );
        
        %sector-level excess supply value
        ESS_s = 1 - (ED_s + X_s)./Y_s ;
        err_ESS_step = max(abs(ESS_s));
        
        %Iterate on domestic sector-level prices
        if err_ESS_step > tol
           p_s_last = p_s;
           p_s = p_s_last - adj_inner*ESS_s;
        else
           break 
        end
    end
    
    %Compute sector-level import price index
    PM_s_next = zeta_s.*( D_s.^( gammaM*omega_star/(1+gammaM*omega_star*kappa) ) );
    PM_s_next(zeta_s==0)=1;
    dif_PM_s = PM_s - PM_s_next ;
    
    err_PM_step = max(abs(dif_PM_s));
    
    %Adjust sector-level import price index
    if err_PM_step > tol
        D_s_next = D_s_last - adj_outter*(D_s_last - D_s);
        
        PM_s = zeta_s.*( D_s_next.^( omega_star/(1+omega_star*kappa) ) );
        PM_s(zeta_s==0)=1;
        D_s_last = D_s_next;
    else
        break 
    end

end

if isnan(err_PM_step) == 1 || isinf(err_PM_step) == 1 || i == 10000
   error('Algorithm did not converge') 
end

%% Compute trade outcomes

%sector-level imports (inclusive of tariffs)
M_s = E_s.*AM_s.*( P_s.^(kappa-1) ).*( PM_s.^(-kappa) );

%product-level imports (inclusive of tariffs)
M_shifter_s = M_s.*( PM_s.^eta );
M_shifter_gs = (M_shifter_s'*Dsg)';

m_g =  ( (   aM_g.*M_shifter_gs                        ).*( varphi_g.^( -eta*(1+gammaM*omega_star*sigma)/(1-sigma) ) ) ).^( 1/(1+gammaM*omega_star*eta) ) ; 
pM_g = ( ( ( aM_g.*M_shifter_gs ).^(gammaM*omega_star) ).*( varphi_g.^(      (1+gammaM*omega_star*sigma)/(1-sigma) ) ) ).^( 1/(1+gammaM*omega_star*eta) ) ; 
pM_g(m_g0 == 0) = 1;

%variety-level imports
M_shifter_g = m_g.*( pM_g.^sigma );
M_shifter_ig = repmat(M_shifter_g', N, 1);

m_ig =      ( (  M_shifter_ig.*a_ig              ).*( z_star_ig.^(-gammaM*sigma) ).*( (1+tau_ig).^(-gammaM*sigma)            ) ).^( 1/(1+gammaM*omega_star*sigma) );
p_star_ig = ( ( (M_shifter_ig.*a_ig).^omega_star ).*( z_star_ig                  ).*( (1+tau_ig).^(-gammaM*sigma*omega_star) ) ).^( 1/(1+gammaM*omega_star*sigma) );

%wlog assign arbitrary values for varieties with zero initial imports
aux = 1./(1+tau_ig);
p_star_ig(m_ig_0==0) = aux(m_ig_0==0); %this is a slow adjustment. I can drop wlog
m_ig(m_ig_0==0) = 0;
pM_ig = (1+tau_ig).*p_star_ig;

%Import values (exclusive tariffs)
VM_ig = m_ig.*p_star_ig;
VM_g = sum(VM_ig,1)';
VM_s = Dsg*VM_g;

%variety-level exports
p_igs = repmat( (p_s'*Dsg) , N, 1);
p_ig = delta_ig.*p_igs;

x_ig = a_star_ig.*( (1+tau_star_ig).*p_ig ).^(-sigma_star);
VX_ig = x_ig.*p_ig;

pX_ig = p_ig.^(gammaX) ; 

%product-level exports
x_g = sum( x_ig,1)';
VX_g = sum(VX_ig,1)';
pX_g = (p_s'*Dsg)';
%pX_g = VX_g./x_g;
%pX_g(x_g==0)=1;

%sector-level exports
VX_s = Dsg*VX_g;
pX_s = (Dsg*pX_g)./sum(Dsg,2);
x_s = VX_s./pX_s;

%sector-level import revenue
RevM_ig = tau_ig.*VM_ig;
RevM_g = sum(RevM_ig, 1)';
RevM_s = Dsg*RevM_g;

%clear z_star_ig a_ig w_rs w_rs_regionterm x_ig pX_ig aux pX_igs M_shifter_ig p_star_ig pM_ig p_ig

%sector-region output/value-added
Y_rs_regionterm = ( Z_rs.*( w_rs.^(-alphaL_rs) ) ).^( 1./alphaK_rs ) ;
Y_rs_sectorterm = ( p_s.*( phi_s.^(-alphaI_s) ) ).^( 1./alphaK_s );
Y_rs = Y_rs_regionterm*diag(Y_rs_sectorterm);
WT_r = Y_rs*(1-alphaI_s);

%Final spending by region
T = D + RevX + sum(RevM_ig, 'all');
T_r = br*T;
F_r = (WT_r + T_r)/betaT;

%Real spending
PT = exp( beta_s'*log( P_s ) );
RW_r = (F_r/PT).^betaT;


%% Outcomes to export

VM_ig = reshape( (m_ig.*pM_ig)', N*G,1);
VM_ig = VM_ig(ind_Msample_ig_vec==1);

pM_ig = reshape(pM_ig', N*G,1);
pM_ig = pM_ig(ind_Msample_ig_vec==1);

RevM_ig = reshape(RevM_ig', N*G,1);
RevM_ig = RevM_ig(ind_Msample_ig_vec==1);

pX_ig = reshape(pX_ig', N*G,1);
pX_ig = pX_ig(ind_Xsample_ig_vec==1);


end

